#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
ONSSaver - make your old savedatas available with new scripts

Copyright (c) 2011-2012 Skydark. All rights reserved.

skydark2 # gmail dot com

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
'''

__author__ = 'Skydark Chen'
__version__ = '0.02.dev'
__license__ = 'GPLv2'

import os
import sys

ENCODING = 'gb18030'
DEBUG_MODE = False


class ONSSaver:
    MAX_SPRITE_NUM = 1000
    MAX_SPRITE2_NUM = 256
    MAX_PARAM_NUM = 100
    MAX_EFFECT_NUM = 256
    
    ONS_NEST_SEARCH_AREA_LIMIT = 50

    def __init__(self, loadhandler, savehandler=None):
        self.loadhandler = loadhandler
        self.savehandler = savehandler
        self.code = ''
        self.length = 0
        self.p = 0

    def reset(self):
        self.p = 0
        self.stack_offset = None
        self.num_nest = 0
        self.nest_info = []
        self.current_offset = (None, )

    def readInt(self, count=1):
        ret = []
        p, code = self.p, self.code
        for i in range(count):
            n = ord(code[p]) + (ord(code[p+1])<<8) + \
                (ord(code[p+2])<<16) + (ord(code[p+3])<<24)
            if ord(code[p+3]) > 8: n -= (2<<31)
            ret.append(n)
            p += 4
        self.p = p
        return ret if len(ret) != 1 else ret[0]

    def readChar(self, count=1):
        cnt = 0
        ret = []
        code, l, p = self.code, self.length, self.p
        while cnt < count and p < l:
            ret.append(code[p])
            cnt += 1
            p += 1
        self.p = p
        return ''.join(ret)

    def readStr(self, count=1):
        ret = []
        code = self.code
        for i in range(count):
            pos = code.find('\x00', self.p)
            assert pos >= 0
            ret.append(''.join(code[self.p:pos]))
            self.p = pos+1
        return ret if len(ret) != 1 else ret[0]

    def readBool(self, count=1):
        b = self.readInt(count)
        if isinstance(b, list):
            assert all(x in (0, 1) for x in b)
        else:
            assert b in (0, 1)
        return b

    def load(self, savedata):
        self.reset()
        self.code = savedata
        self.length = len(savedata)
        readChar = self.readChar
        assert readChar(3) == 'ONS'
        file_version = ord(readChar())*100+ord(readChar())
        printd('Save version: %s' %(file_version))
        self.doPreCheck(file_version)
        self.getStackOffset()
        self.doPostCheck(file_version)
        self.getCurrentLine()

    def doPreCheck(self, file_version):
        readInt, readChar, readStr, readBool =\
            self.readInt, self.readChar, self.readStr, self.readBool
        assert readInt() == 1
        readBool(), readBool() #bold, shadow
        assert readInt() == 0
        readBool() #rmode
        readInt(3) #font rgb
        readStr(2) #cursor
        readInt(2), readStr() #window effect, effect_image
        readInt(8) #font
        readChar(3) #window color
        readChar() #font transparent
        readInt() #waittime
        readInt(4), readStr() #animation info, animation image
        readInt(2) #absflag
        readInt(4) #cursor info
        readStr() #bgimage
        readStr(3), readInt(3), readInt(3) #tachi image, tachi x, tachi y
        assert readInt(3) == [0]*3
        if file_version >= 203: assert readInt(3) == [-1]*3
        for i in range(self.MAX_SPRITE_NUM):
            #ai image, xy, visible, cell, trans
            sprite_info = [readStr(), readInt(), readInt(), readBool(), readInt()]
            if file_version >= 203: sprite_info.append(readInt()) #trans
            #if sprite_info[0]: print i, sprite_info
        self.readVariables(0, self.loadhandler.global_variable_border)

    def doPostCheck(self, file_version):
        readInt, readChar, readStr, readBool =\
            self.readInt, self.readChar, self.readStr, self.readBool
        readBool(), readInt(3) # monocro & color
        readInt() #nega mode
        readStr(2) #midi, wave
        readInt() #cdtrack
        readBool(5) #playloop, waveloop, playonce, bgmloop, mp3loop
        readStr() #musicname
        readBool() #erase text window
        assert readInt() == 1
        for i in range(self.MAX_PARAM_NUM):
            j = readInt()
            if j:
                readInt(5) #ai:x,y,max_w,h,max_param
                readChar(3) #color
                assert ord(readChar()) == 0
            else:
                assert readInt(6) == [-1,0,0,0,0,0]
        for i in range(self.MAX_PARAM_NUM):
            j = readInt()
            readInt(5) #FIXME
        assert readInt(3) == [1, 0, 1]
        readStr() #btndef image
        if file_version >= 202:
            self.readArrayVariable() #<TODO>read array variables
        #<TODO>
        return

    def readArrayVariable(self):
        pass

    def getStackOffset(self):
        readInt = self.readInt
        num_nest = readInt()
        nest_info = []
        self.stack_offset = self.p
        if num_nest > 0:
            ons = self.loadhandler
            printd('nested info: %s' %num_nest)
            self.p += (num_nest-1)*4
            while num_nest > 0:
                i = readInt();
                if i > 0:
                    nest_info.append(('label', i))
                    self.p -= 8
                    num_nest -= 1
                else:
                    self.p -= 16
                    #info_var_no, info_to, info_step = readInt(3)
                    nest_info.append((('for', readInt(3)), -i, ))
                    self.p -= 16
                    num_nest -= 4;
            num_nest = readInt()
            self.p += num_nest*4
            self.num_nest = num_nest
            self.nest_info = [0]*len(nest_info)
            nest_info.reverse()
            for j, (type_, addr) in enumerate(nest_info):
                label = ons.getLabelByAddress(addr)
                line = ons.getLineByAddress(addr, relative=False)
                line_offset = line - label.start_line
                line_start_addr = ons.getAddressByLine(line)
                self.nest_info[j] = (type_, addr, label, line, line_offset, line_start_addr)
                #try:
                #    print type_, addr, \
                #            ons.buf[addr:ons.buf.find('\n', addr+1)].decode(ENCODING)
                #except UnicodeDecodeError:
                #    print type_, addr
                #print label.name

    def getCurrentLine(self):
        p = self.length - 1
        code = self.code
        loadhandler = self.loadhandler
        if code[p] == '*':
            assert code[p-1] == '"'
            p -= 2
            while code[p] != '"': p -= 1
            p -= 1
        self.p = p - 7
        line, current_inline_offset = self.readInt(2)
        label = loadhandler.getLabelByLine(line)
        line_offset = line - label.start_line
        line_start_addr = loadhandler.getAddressByLine(line)
        addr = line_start_addr
        for i in range(current_inline_offset):
            while loadhandler.buf[addr] != ':': addr += 1
            addr += 1
        self.current_offset = (
            self.p, current_inline_offset, 
            addr, label, line, line_offset, line_start_addr
            )
        printd('savestr: %s' %(code[self.p:].decode(ENCODING)))

    def readVariables(self, from_, to_):
        printd('Variable size: %s' %to_)
        for i in range(from_, to_):
            k, s = self.readInt(), self.readStr()
            #if k != 0: print 'Integer variable #%s: %s', i, k
            #if s: print 'String variable #%s: %s', i, s

    def writeInt(self, i):
        if i < 0: i += (2<<31)
        a, b = i % 256, i / 256
        b, c = b % 256, b / 256
        c, d = c % 256, c / 256
        assert 0 <= d < 256
        return chr(a)+chr(b)+chr(c)+chr(d)

    def save(self, savehandler=None):
        loadhandler = self.loadhandler
        if savehandler is None: savehandler = self.savehandler
        if savehandler.global_variable_border != loadhandler.global_variable_border:
            raise ONScripterScriptError('New script has a different global variable border!')
        stack_offset_begin = self.stack_offset
        stack_offset_end = stack_offset_begin + self.num_nest*4
        current_offset_end = self.current_offset[0] 
        current_offset_begin = current_offset_end - 8
        code = self.code

        plain_nest_info = ''
        for type_, addr in self._match_nest_info(savehandler):
            if type_ == 'label':
                plain_nest_info += self.writeInt(addr)
                #try:
                #    print addr, savehandler.buf[addr:addr+20].decode('gbk')
                #except UnicodeDecodeError:
                #    print addr
            else:
                plain_nest_info += ''.join(self.writeInt(i) for i in type_[1])
                plain_nest_info += self.writeInt(-addr)

        current_offset = self._match_current_offset(savehandler)

        ret = '%s%s%s%s%s' %(
            code[:stack_offset_begin],
            plain_nest_info,
            code[stack_offset_end:current_offset_begin],
            current_offset,
            code[current_offset_end:]
            )
        return ret

    def _match_current_offset(self, savehandler):
        p, current_inline_offset = self.current_offset[:2]
        address = self._match_offset(savehandler, *self.current_offset[2:])
        current_lno = savehandler.getLineByAddress(address, relative=False)
        return self.writeInt(current_lno) + self.writeInt(current_inline_offset)

    def _match_nest_info(self, savehandler):
        new_nest_info = [0]*len(self.nest_info)
        for j, nest_info in enumerate(self.nest_info):
            type_ = nest_info[0]
            new_nest_info[j] = (type_, self._match_offset(savehandler, *nest_info[1:]))
        return new_nest_info

    def _match_offset(self, savehandler, addr, label, line, line_offset, line_start_addr):
        # <TODO: other match methods, like diff>
        line_content = self.loadhandler.getLineStringByAddress(line_start_addr)
        try:
            new_label = savehandler.findLabel(label.name)[1]
        except ONSLabelNotFound:
            # Label not found, can't match
            raise ONScripterNoNestMatchFound(str(addr))
        # Label found, try to match with the same line offset
        def try_this_line(newline):
            new_addr = savehandler.getAddressByLine(newline)
            if savehandler.getLineStringByAddress(new_addr) == line_content:
                # matched!
                return new_addr+addr-line_start_addr
            return None
        c = 0
        new_line = new_label.start_line + line_offset
        ret = try_this_line(new_line)
        if ret is not None: return ret
        while c < self.ONS_NEST_SEARCH_AREA_LIMIT:
            c += 1
            ret = try_this_line(new_line+c)
            if ret is not None: return ret
            ret = try_this_line(new_line-c)
            if ret is not None: return ret
        # no match found 
        raise ONScripterNoNestMatchFound(str(addr))


class ONScripterLabel:
    name = ''
    start_line = -1
    start_address = -1
    def __init__(self, name, start_address, start_line):
        self.name = name
        self.start_address = start_address
        self.start_line = start_line


class ONSHandler:
    def __init__(self, directory):
        self.global_variable_border = 200
        self.dir_ = directory
        self.buf = ''
        self.label_info = []
        change_work_directory_once(directory, self.read)

    def read(self):
        buf = ''
        for start in ('0.txt', '00.txt', 'nscript.dat'):
            try:
                buf = open(start, 'rb').read()
                break
            except IOError:
                pass
        else:
            raise ONScripterNoScriptFound('No script found!')
        if start == 'nscript.dat':
            buf = ''.join(chr(ord(c)^0x84) for c in buf)
        else:
            bufs = [buf]
            for i in range(1, 100):
                try:
                    text = open('%d.txt'%(i), 'rb').read()
                except IOError:
                    if i >= 10: continue
                    try:
                        text = open('%02d.txt'%(i), 'rb').read()
                    except IOError:
                        continue
                bufs.append(text)
            buf = '\n'.join(bufs)
        buf = buf.replace('\r\n', '\n').replace('\r', '\n')
        p = 1
        while buf[0] == ';':
            if buf[p:p+4] == 'mode':
                p += 7
            elif buf[p:p+5] == 'value':
                p += 5
                while buf[p] in ('\t', ' ', '\n'): p += 1
                q = p
                while '0' <= buf[p] <= '9': p += 1
                self.global_variable_border = int(buf[q:p])
            else:
                break
            if buf[p] != ',': break
            p += 1
        self.buf = buf
        #open('resultbuf.txt', 'wb').write(buf)
        self.readLabel()

    def readLabel(self):
        state, label, lno = 'newline', '', 0
        label_info = []
        for i, c in enumerate(self.buf):
            if state == 'label':
                c = c.upper()
                if 'A' <= c <= 'Z' or c == '_' or '0' <= c <= '9':
                    label += c
                else:
                    assert self.buf[i-len(label)-1] == '*'
                    state = 'newline' if c == '\n' else 'ready'
                    label_info.append(ONScripterLabel(label, i-len(label)-1, lno))
                    label = ''
                    if c == '\n': lno += 1
                continue
            elif state == 'newline':
                if c in ('\t', ' '): continue
                if c == '*':
                    state, label = 'label', ''
                    continue
                state = 'ready'
            if c == '\n':
                state, lno = 'newline', lno+1
        self.label_info = label_info

    def getLineStringByAddress(self, addr):
        return self.buf[addr:self.buf.find('\n', addr)]

    def getLabelByLine(self, line):
        for i, label in enumerate(self.label_info):
            if label.start_line > line: return self.label_info[i-1]
        return self.label_info[-1]

    def getLabelByAddress(self, address):
        for i, label in enumerate(self.label_info):
            if label.start_address > address: return self.label_info[i-1]
        return self.label_info[-1]

    def getAddressByLine(self, line):
        label = self.getLabelByLine(line)
        l = line - label.start_line
        addr = label.start_address
        buf = self.buf
        length = len(buf)
        while l > 0:
            while addr < length and buf[addr] != '\n': addr += 1
            addr += 1
            l -= 1
        return addr

    def getLineByAddress(self, address, relative=True):
        label = self.getLabelByAddress(address)
        addr = label.start_address
        line = 0 if relative else label.start_line
        buf = self.buf
        while address > addr:
            if buf[addr] == '\n': line += 1
            addr += 1
        return line

    def findLabel(self, name):
        name = name.upper()
        for i, label in enumerate(self.label_info):
            if label.name == name: return i, label
        raise ONSLabelNotFound('No label %s found!' %name)


class ONScripterHandlerError(Exception):
    pass
class ONSLabelNotFound(ONScripterHandlerError):
    pass
class ONScripterNoNestMatchFound(ONScripterHandlerError):
    pass
class ONScripterNoScriptFound(ONScripterHandlerError):
    pass
class ONScripterScriptError(ONScripterHandlerError):
    pass


def printd(s):
    if DEBUG_MODE: print s

def change_work_directory_once(directory, func, *args, **kwargs):
    curpath = os.path.abspath('.')
    try:
        os.chdir(directory)
    except (IOError, OSError), e:
        raise ONScripterNoScriptFound(e)
    ret = func(*args, **kwargs)
    os.chdir(curpath)
    return ret

def help():
    print 'Usage: %s old_script_dir new_script_dir [save_files]' %sys.argv[0]
    print '       (save_files default to all save files in old_script_dir)'

def main(argv):
    if len(argv) < 3:
        help()
        sys.exit()
    #oldsave, newsave, olddir, newdir = sys.argv[1:5]
    load_dir, save_dir = argv[1:3]
    save_files = argv[3:]
    if not save_files:
        import glob
        save_files = glob.glob(os.path.join(load_dir, 'save*.dat'))

    try:
        loadhandler = ONSHandler(load_dir)
        savehandler = ONSHandler(save_dir)
    except ONScripterHandlerError, e:
        print 'Error! Can not load scripts!', e
        return False

    saver = ONSSaver(loadhandler, savehandler)

    for save_file in save_files:
        print 'Processing %s...' %save_file,
        try:
            savedata = open(save_file, 'rb').read()
        except IOError, e:
            print 'Error!\n\t',
            print 'Can not read %s!' %save_file, e
            continue
        try:
            saver.load(savedata)
        except ONScripterHandlerError, e:
            print 'Can not load %s!' %save_file, e
            continue
        except AssertionError, e:
            print 'Error!\n\t',
            print 'Not an onscripter save file? %s' %save_file, e
            continue
        try:
            new_savedata = saver.save()
        except ONScripterHandlerError, e:
            print 'Error!\n\t',
            print 'Can not save %s!' %save_file, e
            continue
        new_save_file = os.path.join(save_dir, os.path.split(save_file)[1])
        try:
            open(new_save_file, 'wb').write(new_savedata)
        except IOError, e:
            print 'Error!\n\t',
            print 'Can not write %s!' %save_file, e
            continue
        print 'Done.'
    print 'OK!'


if __name__=='__main__':
    main(sys.argv)

